package implementsstatic;

import java.io.IOException;
import java.io.Writer;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Filer;
import javax.annotation.processing.Messager;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedAnnotationTypes;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.element.AnnotationValue;
import javax.lang.model.element.Element;
import javax.lang.model.element.ElementKind;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.Modifier;
import javax.lang.model.element.TypeElement;
import javax.lang.model.element.TypeParameterElement;
import javax.lang.model.element.VariableElement;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.ExecutableType;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.ElementFilter;
import javax.lang.model.util.Elements;
import javax.lang.model.util.Types;
import javax.tools.Diagnostic.Kind;


@SupportedSourceVersion(SourceVersion.RELEASE_6)
@SupportedAnnotationTypes("implementsstatic.ImplementsStatic")
public class ImplementsStaticProcessor extends AbstractProcessor {
	
	private static final String IMPLEMENTS_STATIC =
			'@' + ImplementsStatic.class.getSimpleName();
	private static final String INTERFACES = "value";
	private static final String ACTUAL_TPARAMS = "actualTypeParameters";


	private Types typeUtils = null;
	private Elements elementUtils = null;
	private TypeMirror annotationType = null; 
	private Messager msgr = null;
	private ErrorReporter reporter = null;


	@Override
	public void init(ProcessingEnvironment processingEnv){
		super.init(processingEnv);
		typeUtils = processingEnv.getTypeUtils();
		elementUtils = processingEnv.getElementUtils();
		annotationType = elementUtils.getTypeElement(
				ImplementsStatic.class.getName() ).asType();
		msgr = processingEnv.getMessager();
	}


	@Override
	public boolean process(Set<? extends TypeElement> annotations,
			RoundEnvironment roundEnv) {

		assert annotations.size() == 1;

		// process each annotated element
		for(TypeElement a: annotations){
			assert typeUtils.isSameType(a.asType(), annotationType);
			for(Element e: roundEnv.getElementsAnnotatedWith(a))
				processAnnotationOf(e);
		}

		return true;
	}


	private void processAnnotationOf(Element e) {
		// check that the annotation is used on a class
		if(!usedOnClass(e))
			return;

		// locate the correct annotation
		AnnotationMirror am = locateAnnotationOf(e);

		this.reporter = new ErrorReporter(msgr, e, am);

		// reconstruct the full types of interfaces
		// (inject the type parameters)
		List<DeclaredType> interfaces = reconstructInterfaces((TypeElement)e, am);
		if(interfaces == null)
			return;

		// filter out public static methods of the annotated class
		List<ExecutableElementType> methods = getPublicStaticMethodsOf(e);

		// Collect only those static methods that actually implement some
		// method from an interface. Also checks that the annotated class
		// statically implements each interface.
		Set<ExecutableElement> matchingMethods = new HashSet<ExecutableElement>();
		boolean pass = true;
		for(DeclaredType t: interfaces){
			Collection<ExecutableElement> matchingInT = filterImplementingMethods(methods, t);
			if(matchingInT != null)
				matchingMethods.addAll(matchingInT);
			else
				pass = false;
		}
		if(!pass)
			return;

		try {
			generateImplementation(e, processingEnv.getFiler(), interfaces, matchingMethods);
		} catch(IOException ex){
			msgr.printMessage(Kind.ERROR, ex.getMessage());
		}
	}


	/**
	 * Returns a subset of {@code methods} that implement some method from
	 * the given interface {@code t}. If not all methods of {@code t} are
	 * implemented by some method in {@code methods}, error will be reported
	 * and {@code null} returned.
	 * @param methods static methods among which the implementations of
	 *   {@code t}'s methods will be searched
	 * @param t interface that shall be statically implemented by methods
	 *   in {@code methods}.
	 * @return A collection of methods that implement interface {@code t}. If
	 *   not all methods of {@code t} are implemented, returns {@code null}.
	 */
	private Collection<ExecutableElement> filterImplementingMethods(
			List<ExecutableElementType> methods, DeclaredType t) {
		boolean success = true;
		List<ExecutableElement> matching = new LinkedList<ExecutableElement>();
		for(ExecutableElement m: ElementFilter.methodsIn(t.asElement().getEnclosedElements())){
			ExecutableElementType signatureMatch = null;
			ExecutableType mt = (ExecutableType)typeUtils.asMemberOf(t, m);
			for(ExecutableElementType sm: methods)
				if(typeUtils.isSubsignature(mt, sm.getType())){
					signatureMatch = sm;
					break;
				}
			if(signatureMatch == null){
				String msg = "No static implementation of method '"
						+ m + "' declared in interface " + t;
				reporter.report(msg);
				success = false;
				continue;
			}
			if(typeUtils.isSubtype(signatureMatch.getType().getReturnType(), mt.getReturnType())){
					matching.add(signatureMatch.getElement());
			} else {
				String msg = "Incompatible return type of method " + m + " declared in " + t;
				reporter.report(msg);
				success = false;
			}
		}
		return success ? matching : null;
	}


	private List<ExecutableElementType> getPublicStaticMethodsOf(
			Element e) {
		List<ExecutableElementType> methods = new LinkedList<ExecutableElementType>();
		for(ExecutableElement m: ElementFilter.methodsIn(e.getEnclosedElements())){
			Set<Modifier> mods = m.getModifiers();
			if(mods.contains(Modifier.STATIC) && mods.contains(Modifier.PUBLIC))
				methods.add(new ExecutableElementType(m, (ExecutableType)typeUtils.asMemberOf((DeclaredType)e.asType(), m)));
		}
		return methods;
	}


	private AnnotationMirror locateAnnotationOf(Element e) {
		AnnotationMirror am = null;
		for(AnnotationMirror m: e.getAnnotationMirrors())
			if(typeUtils.isSameType(m.getAnnotationType(), annotationType)){
				am = m;
				break;
			}
		assert am != null;
		return am;
	}


	private boolean usedOnClass(Element e) {
		if(!e.getKind().isClass()){
			String msg = e.getKind().isInterface() ?
					IMPLEMENTS_STATIC + " is not allowed on interfaces" :
						IMPLEMENTS_STATIC + " is only allowed on classes";
			msgr.printMessage(Kind.ERROR, msg, e);
			return false;
		} else
			return true;
	}


	private void generateImplementation(Element cls, Filer filer, List<DeclaredType> interfaces, Collection<ExecutableElement> methods) throws IOException {
		assert !interfaces.isEmpty();

		String simpleName = StaticImplementation.getImplementationClassName(cls.getSimpleName().toString());
		String fullName = StaticImplementation.getImplementationClassName(cls.toString());
		String packageName = elementUtils.getPackageOf(cls).toString();

		Writer w = filer.createSourceFile(fullName).openWriter();

		w.append("package ").append(packageName).append(";\n\n");

		w.append("public class ").append(simpleName).append(" implements ");

		Iterator<DeclaredType> i = interfaces.iterator();
		w.append(i.next().toString());
		while(i.hasNext())
			w.append(", ").append(i.next().toString());

		w.append(" {\n\n");

		for(ExecutableElement m: methods)
			generateMethod(m, cls.getSimpleName().toString(), w);
		
		w.append('}');

		w.close();
	}

	private void generateMethod(ExecutableElement m, String className, Writer w) throws IOException {

		w.append("\tpublic ");

		// type parameters
		printTypeParameters(m, w);

		// return type and method name
		w.append(m.getReturnType().toString()).append(' ').append(m.getSimpleName());

		// method parameters
		w.append('(');
		printMethodParameters(m, w);
		w.append(") ");

		// throws clause
		printThrowsClause(m, w);

		// method body begin
		w.append("{\n");
		w.append("\t\t");

		// return statement if necessary
		if(!m.getReturnType().getKind().equals(TypeKind.VOID))
			w.append("return ");

		// static method call begin
		w.append(className).append('.').append(m.getSimpleName());

		// arguments
		w.append('(');
		printMethodArguments(m, w);
		w.append(')');

		// static method call end
		w.append(";\n");

		// method body end
		w.append("\t}\n\n");
	}


	private void printMethodArguments(ExecutableElement m, Writer w) throws IOException {
		List<? extends VariableElement> names = m.getParameters();
		if(!names.isEmpty()){
			Iterator<? extends VariableElement> ni = names.iterator();
			w.append(ni.next().getSimpleName());
			while(ni.hasNext())
				w.append(", ").append(ni.next().getSimpleName());
		}
	}


	private void printThrowsClause(ExecutableElement m, Writer w) throws IOException {
		List<? extends TypeMirror> throwables = m.getThrownTypes();
		if(!throwables.isEmpty()){
			Iterator<? extends TypeMirror> j = throwables.iterator();
			w.append("throws ").append(j.next().toString());
			while(j.hasNext())
				w.append(", ").append(j.next().toString());
			w.append(' ');
		}
	}


	private void printMethodParameters(ExecutableElement m, Writer w) throws IOException {
		List<? extends VariableElement> names = m.getParameters();
		List<? extends TypeMirror> types = ((ExecutableType)m.asType()).getParameterTypes();
		assert names.size() == types.size();
		if(!names.isEmpty()){
			Iterator<? extends VariableElement> ni = names.iterator();
			Iterator<? extends TypeMirror> ti = types.iterator();
			w.append(ti.next().toString()).append(' ').append(ni.next().getSimpleName());
			while(ni.hasNext())
				w.append(", ").append(ti.next().toString()).append(' ').append(ni.next().getSimpleName());
		}
	}


	private void printTypeParameters(ExecutableElement m, Writer w) throws IOException {
		List<? extends TypeParameterElement> tp = m.getTypeParameters();
		if(!tp.isEmpty()){
			w.append('<');
			Iterator<? extends TypeParameterElement> tpi = tp.iterator();
			appendTypeParameter(w, tpi.next());
			while(tpi.hasNext()){
				w.append(", ");
				appendTypeParameter(w, tpi.next());
			}
			w.append("> ");
		}
	}


	private void appendTypeParameter(Writer to, TypeParameterElement tpe) throws IOException {
		to.append(tpe.getSimpleName());
		List<? extends TypeMirror> bounds = tpe.getBounds();
		if(!bounds.isEmpty()){
			Iterator<? extends TypeMirror> bi = bounds.iterator();
			to.append(" extends ").append(bi.next().toString());
			while(bi.hasNext())
				to.append(" & ").append(bi.next().toString());
		}
	}

	private List<DeclaredType> reconstructInterfaces(TypeElement e, AnnotationMirror am){

		// get the values of annotation parameters
		List<DeclaredType> interfaces = null;
		List<DeclaredType> aparams = null;
		for(Map.Entry<? extends ExecutableElement, ? extends AnnotationValue> entry: elementUtils.getElementValuesWithDefaults(am).entrySet()){
			String eename = entry.getKey().getSimpleName().toString();
			if(eename.equals(INTERFACES))
				interfaces = unwrap((List<AnnotationValue>)entry.getValue().getValue(), DeclaredType.class);
			else if(eename.equals(ACTUAL_TPARAMS))
				aparams = unwrap((List<AnnotationValue>)entry.getValue().getValue(), DeclaredType.class);
		}

		// check that the list of interfaces is non-empty
		if(interfaces.isEmpty()){
			msgr.printMessage(Kind.ERROR, "Empty array", e, am);
			return null;
		}

		// check that provided classes are all interfaces
		if(!areAllInterfaces(interfaces))
			return null;

		// go through the interfaces and reconstruct their type parameters
		Iterator<DeclaredType> ai = aparams.iterator();
		List<DeclaredType> reconstructed = new ArrayList<DeclaredType>(interfaces.size());
		for(DeclaredType t: interfaces){
			DeclaredType r = reconstructType(t, e, ai);
			if(r == null)
				return null;
			reconstructed.add(r);
		}

		// check that there are no type parameters left
		if(ai.hasNext()){
			String msg = "More type parameters than required";
			msgr.printMessage(Kind.ERROR, msg, e, am);
		}

		return reconstructed;
	}


	private DeclaredType reconstructType(DeclaredType t, TypeElement e,
			Iterator<DeclaredType> ai) {

		TypeElement te = (TypeElement)t.asElement();
		List<? extends TypeParameterElement> params = te.getTypeParameters();
		if(params.isEmpty())
			return t;

		int n = params.size();
		TypeMirror[] args = new TypeMirror[n];
		for(int i=0; i<n; ++i){
			if(!ai.hasNext()){
				reporter.report("Not enough type arguments");
				return null;
			}
			DeclaredType a = ai.next();
			args[i] = reconstructType(a, e, ai);
			if(args[i] == null)
				return null;
		}
		
		return typeUtils.getDeclaredType(te, args);
	}

	private boolean areAllInterfaces(List<DeclaredType> typeList) {
		boolean res = true;
		for(DeclaredType t: typeList){
			if(!t.asElement().getKind().equals(ElementKind.INTERFACE)){
				reporter.report(t + " is not an interface");
				res = false;
			}
		}
		return res;
	}


	private <T> List<T> unwrap(List<AnnotationValue> srcList, Class<T> tgtType) {
		List<T> l = new ArrayList<T>(srcList.size());
		for(AnnotationValue v: srcList)
			l.add(tgtType.cast(v.getValue()));
		return l;
	}
}


/**
 * A pair of {@link ExecutableElement} and {@link ExecutableType}.
 */
class ExecutableElementType {
	private final ExecutableElement elem;
	private final ExecutableType type;
	public ExecutableElementType(ExecutableElement elem, ExecutableType type){
		this.elem = elem;
		this.type = type;
	}
	public ExecutableElement getElement() {
		return elem;
	}
	public ExecutableType getType() {
		return type;
	}
}


/**
 * A convenience class for reporting errors.
 */
class ErrorReporter {
	private final Messager msgr;
	private final Element elem;
	private final AnnotationMirror am;

	public ErrorReporter(Messager msgr, Element elem, AnnotationMirror am){
		this.msgr = msgr;
		this.elem = elem;
		this.am = am;
	}

	public void report(String msg){
		msgr.printMessage(Kind.ERROR, msg, elem, am);
	}
}